import streamlit as st
import time
from streamlit_option_menu import option_menu
from data.user_data import insert_user_training

# 设置Streamlit应用程序的标题
st.set_page_config(
    page_title="goudaner.AI",
    page_icon="🧊",
    layout="wide",
    initial_sidebar_state="expanded"
)
from ChatbotWithRetrieval import ChatbotWithRetrieval
import os
from customer_logging import get_logger

logger = get_logger("sys-")
# 用户名和密码的默认值
user = {
    "admin": {"password": "admin1",
              "model_type": "yy", "training_table": "user_training_yy"},
    "root": {"password": "root1",
             "model_type": "zeekr", "training_table": "user_training_zeekr"}
}

# 检查会话状态中是否有登录状态，如果没有，初始化为 False
if 'logged_in' not in st.session_state:
    st.session_state.logged_in = False


def login_page():
    with st.form("login_form"):
        st.title("登录")
        username = st.text_input("用户名", value="")
        password = st.text_input("密码", value="", type="password")
        submit = st.form_submit_button("登录")

        if submit:
            user_info = user.get(username)
            user_password = user_info.get("password")
            if user_password is None:
                return st.error("用户名不能存在")
            if password == user_password:
                st.success("登录成功！")
                # 更新会话状态为已登录
                st.session_state.logged_in = True
                model_type = user_info.get("model_type")
                training_table = user_info.get("training_table")
                st.session_state.model_type = model_type
                st.session_state.training_table = training_table
                st.experimental_rerun()  # 重新运行脚本以显示主页面
            else:
                st.error("用户名或密码错误，请重新输入。")


def get_llm():
    if st.session_state.model_type == "zeekr":
        return get_zeekr_llm()
    else:
        return get_yy_llm()


@st.cache_resource
def get_yy_llm():
    logger.info("加载对话大模型")
    base_dir = 'One'
    llm = ChatbotWithRetrieval("my_vectors.faiss")
    return llm


@st.cache_resource
def get_zeekr_llm():
    logger.info("加载对话大模型")
    base_dir = 'One'
    llm = ChatbotWithRetrieval("zeekr_vectors.faiss")
    return llm


# 设置Streamlit应用程序的标题
# st.set_page_config(page_title="goudaner.AI", layout="wide")


## 上传并预览（1M以内才可预览）
# 上传文件

def upload_Save():
    if "uploaded_state" not in st.session_state:
        st.session_state.uploaded_state = False
    if st.session_state.uploaded_state:
        return
    # 创建一个文件夹用于保存上传的文件（若存在则清空，若不存在，则新建）
    dirs = 'One'
    if not os.path.exists(dirs):
        os.makedirs(dirs)
    # 选择文件
    uploaded_files = st.file_uploader("###### 您可以再此处上传知识文件", accept_multiple_files=True,
                                      type=["pdf", "txt", "docx"])
    # 保存文件
    if uploaded_files is None:
        return

    for uploaded_file in uploaded_files:
        file_contents = uploaded_file.getvalue()
        file_path = os.path.join(dirs, uploaded_file.name)
        llm = get_llm()
        fassStore = llm.store
        # 将文件保存到本地文件系统
        with open(file_path, "wb") as f:
            f.write(file_contents)
        # 获取文件路径
        fassStore.load_file(file_path)
        st.session_state.uploaded_state = True
        st.write(f"文件地址: {file_path}")
    return os.path.join(os.path.dirname(os.path.abspath(__file__)), dirs)


def main():
    st.sidebar.title("goudaner.AI v0.1")

    menu1 = "知识库ai"
    menu4 = "添加-知识"

    with st.sidebar:
        menu = option_menu("功能分类", [menu1, menu4],
                           icons=['house', "list-task"],
                           menu_icon="cast", default_index=0)
    if menu == menu1:
        st.title("💬 知识库ai")
        st.caption("🚀 狗蛋1号")
        if st.button('清除聊天记录'):
            st.session_state.history = []
        showLLMChatbot()
    if menu == menu4:
        # 文件上传
        upload_Save()
        knowledge_training_data()


def knowledge_training_data():
    # 定义selectbox来获取数据类型
    data_type = st.selectbox("训练数据类型", ("问答示例", "文档"))
    # 使用st.form来处理提交事件
    with st.form("t_d"):
        training_type = data_type == "问答示例"
        input_type = "问题" if training_type else "标题"
        answer_type = "回答" if training_type else "内容"
        input_info = st.text_input(input_type, max_chars=100, help='最大长度为100字符')
        answer_info = st.text_area(label=answer_type,
                                   value='请输入...',
                                   height=5,
                                   max_chars=300,
                                   help='最大长度限制为300')
        # 提交按钮
        submitted = st.form_submit_button("训练")
        if submitted:
            llm = get_llm()
            faiss = llm.store
            template = f"{input_type} : {input_info} ;   {answer_type} : {answer_info}"
            user_training_data = {
                "type": input_type,
                "content": template,
                "state": "Untrained"
            }
            faiss.write_templates(template)
            training_table = st.session_state.training_table
            insert_user_training(training_table, [user_training_data])


@st.cache_data(show_spinner="正在努力生成答案 ...", ttl=60 * 60)
def get_response_material(user_input, history, model_type):
    llm = get_llm()
    chat_chain = llm.qa

    logger.info("用户对话问题：" + user_input)
    response = chat_chain.invoke({"history": history, "input": user_input})["answer"]
    logger.info("用户对话回答：" + response)
    # 返回两个值
    return response


def showLLMChatbot():
    # 给对话增加history属性，将历史对话信息储存下来
    if "history" not in st.session_state:
        st.session_state.history = []
        assistant_message = st.chat_message(
            "assistant", avatar="https://i.geely.com/favicon.ico"
        )
        assistant_message.markdown("尊敬的用户您好，您的吉利伙伴assistant很高兴为您服务，请问有什么可以帮您？")

    # 显示历史信息
    for message in st.session_state.history:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    user_input = st.chat_input()
    if user_input is None:
        return
    # 在页面上显示用户的输入
    with st.chat_message("user"):
        st.markdown(user_input)

    response = get_response_material(user_input, st.session_state.history, st.session_state.model_type)
    # 在页面上显示模型生成的回复
    with st.chat_message("assistant", avatar="https://i.geely.com/favicon.ico"):
        st.markdown(response)

    # 将用户的输入加入历史
    st.session_state.history.append({"role": "user", "content": user_input})
    # 将模型的输出加入到历史信息中
    st.session_state.history.append({"role": "assistant", "content": response})

    # 只保留十轮对话，这个可根据自己的情况设定，我这里主要是会把history给大模型，context有限，轮数不能太多
    if len(st.session_state.history) > 11:
        st.session_state.history = st.session_state.history[-11:]


if __name__ == '__main__':

    # 如果用户已登录，则显示应用的其他部分
    if st.session_state.logged_in:
        main()
    else:
        # 如果用户未登录，则显示登录页面
        login_page()
